// SPDX-License-Identifier: Mulan PSL v2
/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This software is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

use std::{
    fmt::{Debug, Display},
    io,
    ops::{Deref, DerefMut},
    os::fd::OwnedFd,
};

use rustix::{fs, io::Errno, shm};
use tracing::trace;

use super::{
    flock::{FileLock, FileLockGuard, FileLockOperation},
    mem,
    mmap::MirroredMmap,
};

#[repr(transparent)]
pub struct Name(String);

impl Name {
    pub fn new<S: AsRef<str>>(name: S) -> Self {
        let name = name.as_ref().trim();
        debug_assert!(!name.is_empty(), "Shmem name cannot be empty");

        Self(format!(
            "/{}",
            name.trim_start_matches('/').replace('/', "_")
        ))
    }
}

impl AsRef<str> for Name {
    fn as_ref(&self) -> &str {
        self.0.as_str()
    }
}

impl Debug for Name {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        Debug::fmt(&self.0, f)
    }
}

impl Display for Name {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        Display::fmt(&self.0, f)
    }
}

/// A handle to a POSIX shared memory object, presented as a high-performance,
/// mirrored ring buffer.
///
/// This structure is the primary interface for creating and accessing shared memory.
/// It simplifies inter-process communication (IPC) by providing a ring buffer that
/// can be read from and written to as if it were a single, contiguous block of memory,
/// eliminating the need for manual wrap-around logic in most cases.
#[derive(Debug)]
pub struct Shmem {
    /// The unique name of the POSIX shared memory object (e.g., "/my-shmem").
    name: Name,
    /// The underlying mirrored memory mapping.
    mmap: MirroredMmap<OwnedFd>,
    /// If `true`, this instance is the owner and will unlink the shared memory object on drop.
    owned: bool,
}

impl Shmem {
    /// Creates a new shared memory object with a mirrored mapping.
    ///
    /// This function creates a new shared memory file and will fail if an object
    /// with the same name already exists. The creator is considered the "owner" and
    /// is responsible for unlinking the object upon drop.
    ///
    /// # Arguments
    ///
    /// * `name`: The name for the new POSIX shared memory object (e.g., `"/my-buffer"`).
    /// * `header`: The desired size of the **header region (A)**. This value will also
    ///   be aligned up to the nearest page boundary.
    /// * `data`: The desired size of the **data region (B)**. This value directly
    ///   influences the return value of `data_len()` and will be aligned up to the
    ///   nearest page boundary.
    ///
    /// # Important: File Size Calculation
    ///
    /// Please note that the final **total size of the created file** will be the sum of
    /// the **page-aligned `reserve` size** and the **page-aligned `len` size**.
    ///
    /// ### Example
    ///
    /// Assuming a system page size of 4096 bytes (4KB), a call to
    /// `create("test", 1000, 200)` will result in the following:
    /// 1. `head_len` = `align_up(200)` = `4096` bytes.
    /// 2. `data_len` = `align_up(1000)` = `4096` bytes.
    /// 3. **Final file size** = `4096 + 4096` = `8192` bytes.
    ///
    /// In this example, `header_len()` will return `4096`, and `data_len()` will
    /// also return `4096`.
    pub fn create<S: AsRef<str>>(name: S, header: usize, data: usize) -> io::Result<Self> {
        if name.as_ref().is_empty() {
            return Err(Errno::INVAL.into());
        }
        let name = Name::new(name);

        trace!("[Shmem] Creating shmem '{}'...", name);
        if data == 0 {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "Data size cannot be zero",
            ));
        }
        let fd = shm::open(
            name.as_ref(),
            shm::OFlags::RDWR | shm::OFlags::CREATE | shm::OFlags::EXCL,
            shm::Mode::from(0o600),
        )?;

        let head_len = mem::page_align_up(header);
        let data_len = mem::page_align_up(data);

        let file_len = head_len.checked_add(data_len).ok_or(Errno::OVERFLOW)?;
        fs::ftruncate(&fd, file_len.try_into().map_err(|_| Errno::OVERFLOW)?)?;

        let mmap = MirroredMmap::mmap(fd, head_len, data_len)?;

        let shmem = Self {
            name,
            mmap,
            owned: true,
        };
        trace!("[Shmem] {:#?}", shmem);
        trace!("[Shmem] Shmem '{}' created.", shmem);

        Ok(shmem)
    }

    /// Opens an existing shared memory object with a mirrored mapping.
    ///
    /// This will fail if the object does not exist. The handle returned is not the owner
    /// and will not unlink the object on drop.
    ///
    /// # Arguments
    /// * `name`: The name of the existing POSIX shared memory object.
    /// * `header`: The size of the header region (A). This **must** match the size
    ///   used when the object was created.
    pub fn open<S: AsRef<str>>(name: S, header: usize) -> io::Result<Self> {
        if name.as_ref().is_empty() {
            return Err(Errno::INVAL.into());
        }
        let name = Name::new(name);

        trace!("[Shmem] Opening shmem '{}'...", name);
        let fd = shm::open(name.as_ref(), shm::OFlags::RDWR, shm::Mode::empty())?;

        let file_len = fs::fstat(&fd)?
            .st_size
            .try_into()
            .map_err(|_| Errno::OVERFLOW)?;
        if !mem::is_page_aligned(file_len) {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("File size must be page-aligned (got {} bytes)", file_len),
            ));
        }

        let head_len = mem::page_align_up(header);
        let data_len = file_len.checked_sub(head_len).ok_or(Errno::OVERFLOW)?;
        let mmap = MirroredMmap::mmap(fd, head_len, data_len)?;

        let shmem = Self {
            name: Name::new(name),
            mmap,
            owned: false,
        };
        trace!("[Shmem] {:#?}", shmem);
        trace!("[Shmem] Shmem '{}' opened.", shmem);

        Ok(shmem)
    }

    /// Returns the name of the shared memory object.
    #[inline]
    pub fn name(&self) -> &str {
        self.name.as_ref()
    }

    /// Acquires an exclusive lock on the underlying file descriptor.
    ///
    /// The lock is released when the returned `` is dropped.
    #[inline]
    pub fn flock(&self) -> io::Result<FileLockGuard<'_>> {
        FileLock::lock(&self.mmap, FileLockOperation::LockExclusive)
    }

    /// Returns the underlying memory mapping.
    #[inline]
    pub fn mmap(&self) -> &MirroredMmap<OwnedFd> {
        &self.mmap
    }

    /// Returns `true` if this handle created the shared memory object.
    ///
    /// The owner is responsible for unlinking the object from the system on drop.
    #[inline]
    pub fn is_owner(&self) -> bool {
        self.owned
    }
}

impl Deref for Shmem {
    type Target = [u8];

    fn deref(&self) -> &Self::Target {
        unsafe { std::slice::from_raw_parts(self.mmap().data_ptr(), self.mmap().data_len()) }
    }
}

impl DerefMut for Shmem {
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { std::slice::from_raw_parts_mut(self.mmap().data_ptr(), self.mmap().data_len()) }
    }
}

// SAFETY: It is safe to send a `Shmem` across threads. The raw pointer refers
// to a process-wide shared memory map, so the address is valid in any thread.
// Ownership transfer ensures resources are properly released by `Drop` exactly
// once on the new owning thread.
unsafe impl Send for Shmem {}

// SAFETY: It is safe to share `&Shmem` across threads because this only provides
// shared access to the handle and metadata. The raw pointer `ptr` inside does
// not grant unsynchronized write access. Any mutation to the underlying memory
// must be done through a synchronized mechanism (like the allocator's lock) or
// via `unsafe` code, where the safety burden is on the caller.
unsafe impl Sync for Shmem {}

impl Display for Shmem {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        Display::fmt(&self.name, f)
    }
}

impl Drop for Shmem {
    fn drop(&mut self) {
        if self.owned {
            let _ = shm::unlink(self.name.as_ref());
        }
    }
}

#[cfg(test)]
pub fn unique_name() -> String {
    use std::{
        process,
        sync::{
            atomic::{AtomicUsize, Ordering},
            OnceLock,
        },
    };

    static SEQ: AtomicUsize = AtomicUsize::new(0);
    static EXEC_NAME: OnceLock<String> = OnceLock::new();

    let exec_name = EXEC_NAME.get_or_init(|| {
        use std::ffi::OsStr;

        std::env::current_exe()
            .ok()
            .and_then(|p| p.file_stem().and_then(OsStr::to_str).map(str::to_owned))
            .unwrap_or_else(|| "unknown".into())
    });

    format!(
        "/{}_{}_{}",
        exec_name,
        process::id(),
        SEQ.fetch_add(1, Ordering::Relaxed)
    )
}

#[cfg(test)]
mod tests {
    use std::{thread, time::Duration};

    use super::*;

    #[test]
    fn test_create_errors() {
        let name = unique_name();
        assert!(Shmem::create(&name, 0, 0).is_err());

        let _shmem = Shmem::create(&name, 1024, 1024).unwrap();
        assert!(Shmem::create(&name, 1024, 1024).is_err());
        assert!(Shmem::create(&name, 2048, 1024).is_err());
    }

    #[test]
    fn test_create_and_open() {
        let name = unique_name();
        let size = 4096;

        let shmem = Shmem::create(&name, size, size).unwrap();
        println!("{:#?}", shmem);
        assert_eq!(shmem.name(), name);
        assert_eq!(shmem.mmap().data_len(), size);
        assert!(shmem.is_owner());

        let shmem2 = Shmem::open(&name, size).unwrap();
        println!("{:#?}", shmem2);
        assert_eq!(shmem2.name(), name);
        assert_eq!(shmem2.mmap().data_len(), size);
        assert!(!shmem2.is_owner());
    }

    #[test]
    fn test_create_and_open_threaded() {
        let name = unique_name();
        let size = 1024;

        let _shmem = Shmem::create(&name, size, size).unwrap();
        let child = thread::spawn(move || {
            thread::sleep(Duration::from_millis(100));
            let shmem_child = Shmem::open(&name, size).unwrap();
            assert_eq!(shmem_child.mmap().data_len(), mem::page_align_up(size));
        });

        child.join().unwrap();
    }

    #[test]
    fn test_cleanup() {
        let name = unique_name();
        {
            let _shmem = Shmem::create(&name, 1024, 1024).unwrap();
        }
        assert!(Shmem::open(&name, 0).is_err());
    }
}
